from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 转科相关基础数据 ==========
DEPARTMENTS = [
    ("D01", "肿瘤内科"),
    ("D02", "肿瘤外科"),
    ("D03", "放疗科"),
    ("D04", "血液科"),
    ("D05", "中西医结合科"),
    ("D06", "重症医学科"),
    ("D07", "疼痛科"),
    ("D08", "营养科")
]

TRANSFER_REASONS = [
    "需要专科治疗",
    "病情变化需调整治疗方案",
    "术后恢复",
    "化疗后观察",
    "放疗后护理",
    "多学科会诊后调整",
    "并发症处理"
]

HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 8 if date_obj.day % 2 == 0 else 5


# ========== 核心生成函数 ==========
def generate_transfer_records():
    sql_statements = []
    record_count = 0

    patient_id_counter = 2001
    visit_sn_counter = 2001
    transfer_id_counter = 70001

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1
            patient_id = f"PT{patient_id_counter}"
            patient_id_counter += 1

            visit_sn = f"VIS{visit_sn_counter}"
            visit_sn_counter += 1

            # 随机选择医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 生成转科记录ID
            transfer_id = f"TRANS{transfer_id_counter}"
            transfer_id_counter += 1

            # 随机选择转出和转入科室（确保不同）
            from_dept_code, from_dept_name = random.choice(DEPARTMENTS)
            to_dept_code, to_dept_name = random.choice([d for d in DEPARTMENTS if d[0] != from_dept_code])

            # 时间信息
            transfer_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d"),
                end_date=(datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1))
            ).strftime('%Y-%m-%d %H:%M:%S')

            record_datetime = transfer_datetime
            yy_collection_datetime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            yy_etl_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # 其他字段
            reason = random.choice(TRANSFER_REASONS)

            # 生成yy_record_md5 (模拟MD5)
            yy_record_md5 = f"md5_{random.getrandbits(128):032x}"
            yy_batch_time = day
            yy_record_batch_id = f"BATCH{day}_{i}"

            data = {
                'from_table': 'INPATIENT_TRANSFER',
                'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                'hospital_code': hospital_code,
                'hospital_name': hospital_name,
                'patient_id': patient_id,
                'patient_id_old': f"OLD_{patient_id}",
                'reason': reason,
                'record_datetime': record_datetime,
                'record_status': 1,
                'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'trans_from_dept_code': from_dept_code,
                'trans_from_dept_name': from_dept_name,
                'trans_to_dept_code': to_dept_code,
                'trans_to_dept_name': to_dept_name,
                'transfer_datetime': transfer_datetime,
                'transfer_id': transfer_id,
                'visit_sn': visit_sn,
                'yy_collection_datetime': yy_collection_datetime,
                'yy_record_md5': yy_record_md5,
                'yy_upload_status': 0,
                'yy_etl_time': yy_etl_time,
                'yy_upload_time': None,
                'yy_batch_time': yy_batch_time,
                'yy_record_batch_id': yy_record_batch_id,
                'yy_backfill_time': None,
                'yy_backfill_status': None,
                'branch_code': branch_code,
                'branch_name': branch_name,
                'date_for_partition': transfer_datetime
            }

            sql = _generate_sql('b03_3', data)
            sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_transfer_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b03_3')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b03_3_transfer_records.sql")